import { app } from "../../scripts/app.js";

// 通用训练Loss图表组件
class TrainingLossChart {
    constructor(node) {
        this.node = node;
        this.lossData = [];
        this.maxDataPoints = 100;
        this.chartVisible = true;
        this.legendHitboxes = {};
        this.chartColor = "#8b5cf6"; // 默认紫色，可自定义
        this.chartTitle = "Training Loss";
        
        console.log(`训练Loss图表初始化: ${node.constructor.comfyClass}`);
    }
    
    // 更新loss数据
    updateLoss(step, loss) {
        if (loss !== undefined && loss !== null) {
            this.lossData.push({
                step: step || 0,
                loss: loss,
                timestamp: Date.now()
            });
            
            // 限制数据点数量
            if (this.lossData.length > this.maxDataPoints) {
                this.lossData = this.lossData.slice(-this.maxDataPoints);
            }
        }
    }
    
    // 设置图表颜色
    setChartColor(color) {
        this.chartColor = color;
    }
    
    // 设置图表标题
    setChartTitle(title) {
        this.chartTitle = title;
    }
    
    // 绘制图表
    draw(ctx) {
        if (!this.chartVisible || this.lossData.length < 2) return;
        
        const chartRect = { 
            x: 15, 
            y: 15, 
            width: this.node.size[0] - 30, 
            height: this.node.size[1] - 120
        };
        
        this.drawTitle(ctx);
        this.drawGrid(ctx, chartRect);
        this.drawLossLine(ctx, chartRect);
        this.drawLegend(ctx, chartRect);
    }
    
    // 绘制图表标题
    drawTitle(ctx) {
        ctx.fillStyle = "#fff";
        ctx.font = "bold 14px Arial";
        ctx.textAlign = "center";
        
        let title = this.chartTitle;
        if (this.node.progress && this.node.progress.status) {
            title += ` (${this.node.progress.status})`;
        }
        ctx.fillText(title, this.node.size[0] / 2, 12);
    }
    
    // 绘制图表网格
    drawGrid(ctx, rect) {
        ctx.strokeStyle = "#333";
        ctx.lineWidth = 1;
        ctx.font = "10px Arial";
        ctx.fillStyle = "#888";
        ctx.textAlign = "right";

        // 计算loss范围
        const losses = this.lossData.map(d => d.loss);
        const minLoss = Math.min(...losses);
        const maxLoss = Math.max(...losses);
        const lossRange = maxLoss - minLoss;
        
        // 绘制水平网格线
        for (let i = 0; i <= 4; i++) {
            const p = i / 4;
            const y = rect.y + p * rect.height;
            const value = maxLoss - p * lossRange;

            ctx.beginPath();
            ctx.moveTo(rect.x, y);
            ctx.lineTo(rect.x + rect.width, y);
            ctx.stroke();
            
            ctx.fillText(`${value.toFixed(4)}`, rect.x - 5, y + 3);
        }
        
        // 绘制垂直网格线（步数）
        const steps = this.lossData.map(d => d.step);
        const minStep = Math.min(...steps);
        const maxStep = Math.max(...steps);
        const stepRange = maxStep - minStep;
        
        for (let i = 0; i <= 4; i++) {
            const p = i / 4;
            const x = rect.x + p * rect.width;
            const value = minStep + p * stepRange;

            ctx.beginPath();
            ctx.moveTo(x, rect.y);
            ctx.lineTo(x, rect.y + rect.height);
            ctx.stroke();
            
            ctx.textAlign = "center";
            ctx.fillText(`${Math.round(value)}`, x, rect.y + rect.height + 12);
        }
        ctx.textAlign = "left";
    }
    
    // 绘制Loss曲线
    drawLossLine(ctx, rect) {
        if (this.lossData.length < 2) return;

        const losses = this.lossData.map(d => d.loss);
        const steps = this.lossData.map(d => d.step);
        const minLoss = Math.min(...losses);
        const maxLoss = Math.max(...losses);
        const lossRange = maxLoss - minLoss;
        const minStep = Math.min(...steps);
        const maxStep = Math.max(...steps);
        const stepRange = maxStep - minStep;
        
        const width = rect.width;
        const height = rect.height;
        const xStep = this.lossData.length > 1 ? width / (this.lossData.length - 1) : width;
        
        // 创建渐变填充
        const grad = ctx.createLinearGradient(0, rect.y, 0, rect.y + height);
        grad.addColorStop(0, this.chartColor + "60"); // 半透明
        grad.addColorStop(1, this.chartColor + "00"); // 透明
        ctx.fillStyle = grad;
        
        // 绘制填充区域
        ctx.beginPath();
        ctx.moveTo(rect.x, rect.y + height);
        this.lossData.forEach((d, i) => {
            const x = rect.x + i * xStep;
            const y = rect.y + height - ((d.loss - minLoss) / lossRange) * height;
            ctx.lineTo(x, y);
        });
        ctx.lineTo(rect.x + (this.lossData.length-1) * xStep, rect.y + height);
        ctx.closePath();
        ctx.fill();

        // 绘制线条
        ctx.strokeStyle = this.chartColor;
        ctx.lineWidth = 2;
        ctx.beginPath();
        this.lossData.forEach((d, i) => {
            const x = rect.x + i * xStep;
            const y = rect.y + height - ((d.loss - minLoss) / lossRange) * height;
            if (i === 0) ctx.moveTo(x, y);
            else ctx.lineTo(x, y);
        });
        ctx.stroke();
        
        // 绘制数据点
        ctx.fillStyle = this.chartColor;
        this.lossData.forEach((d, i) => {
            const x = rect.x + i * xStep;
            const y = rect.y + height - ((d.loss - minLoss) / lossRange) * height;
            ctx.beginPath();
            ctx.arc(x, y, 2, 0, 2 * Math.PI);
            ctx.fill();
        });
    }
    
    // 绘制图表图例
    drawLegend(ctx, rect) {
        const legendY = rect.y + rect.height + 25;
        const currentLoss = this.lossData.length > 0 ? this.lossData[this.lossData.length - 1].loss : 0;
        
        ctx.font = "12px Arial";
        let currentX = rect.x;
        this.legendHitboxes = {};

        // Loss图例
        const isVisible = this.chartVisible;
        ctx.globalAlpha = isVisible ? 1.0 : 0.4;

        ctx.fillStyle = this.chartColor;
        ctx.fillRect(currentX, legendY - 8, 10, 10);
        
        ctx.fillStyle = "#fff";
        const text = `Loss: ${currentLoss.toFixed(4)}`;
        ctx.fillText(text, currentX + 15, legendY);
        
        const itemWidth = ctx.measureText(text).width + 20;
        this.legendHitboxes["chart"] = {
            x: currentX,
            y: legendY - 10,
            width: itemWidth,
            height: 20
        };
        
        currentX += itemWidth + 15;
        
        // 数据点数量
        ctx.fillStyle = "#fff";
        ctx.fillText(`数据点: ${this.lossData.length}`, currentX, legendY);
        
        ctx.globalAlpha = 1.0;
    }
    
    // 处理鼠标点击
    handleMouseClick(pos) {
        for (const id in this.legendHitboxes) {
            const box = this.legendHitboxes[id];
            if (pos[0] > box.x && pos[0] < box.x + box.width &&
                pos[1] > box.y && pos[1] < box.y + box.height) {
                
                if (id === "chart") {
                    this.chartVisible = !this.chartVisible;
                    return true;
                }
            }
        }
        return false;
    }
    
    // 清除数据
    clear() {
        this.lossData = [];
    }
    
    // 获取当前loss
    getCurrentLoss() {
        return this.lossData.length > 0 ? this.lossData[this.lossData.length - 1].loss : null;
    }
    
    // 获取loss趋势
    getLossTrend() {
        if (this.lossData.length < 2) return "stable";
        
        const recent = this.lossData.slice(-10); // 最近10个点
        const first = recent[0].loss;
        const last = recent[recent.length - 1].loss;
        
        if (last < first * 0.95) return "decreasing";
        if (last > first * 1.05) return "increasing";
        return "stable";
    }
}

// 注册通用训练图表扩展
app.registerExtension({
    name: "Comfy.TrainingLossChart",
    
    async beforeRegisterNodeDef(nodeType) {
        // 支持所有训练节点
        const trainingNodes = [
            "SimpleFLUXLoRATrainer",
            "AdvancedLoRATrainer", 
            "SDXLLoRATrainer",
            "SD3LoRATrainer"
        ];
        
        if (!trainingNodes.includes(nodeType.comfyClass)) {
            return;
        }
        
        console.log(`注册训练Loss图表: ${nodeType.comfyClass}`);
        
        // 增加节点大小以容纳图表
        nodeType.size = [400, 300];
        
        const onNodeCreated = nodeType.prototype.onNodeCreated;
        nodeType.prototype.onNodeCreated = function() {
            onNodeCreated?.apply(this, arguments);
            
            // 创建Loss图表实例
            this.lossChart = new TrainingLossChart(this);
            
            // 根据节点类型设置颜色和标题
            switch (this.constructor.comfyClass) {
                case "SimpleFLUXLoRATrainer":
                    this.lossChart.setChartColor("#8b5cf6"); // FLUX紫色
                    this.lossChart.setChartTitle("FLUX Training Loss");
                    break;
                case "AdvancedLoRATrainer":
                    this.lossChart.setChartColor("#3b82f6"); // 蓝色
                    this.lossChart.setChartTitle("Advanced LoRA Training Loss");
                    break;
                case "SDXLLoRATrainer":
                    this.lossChart.setChartColor("#10b981"); // 绿色
                    this.lossChart.setChartTitle("SDXL LoRA Training Loss");
                    break;
                case "SD3LoRATrainer":
                    this.lossChart.setChartColor("#f59e0b"); // 橙色
                    this.lossChart.setChartTitle("SD3 LoRA Training Loss");
                    break;
            }
            
            console.log(`${this.constructor.comfyClass} 训练节点初始化完成`);
        };
        
        const onDrawForeground = nodeType.prototype.onDrawForeground;
        nodeType.prototype.onDrawForeground = function(ctx) {
            if (onDrawForeground) {
                onDrawForeground.call(this, ctx);
            }
            
            // 绘制Loss图表
            if (this.lossChart) {
                this.lossChart.draw(ctx);
            }
        };
        
        // 添加鼠标点击事件处理
        const onMouseDown = nodeType.prototype.onMouseDown;
        nodeType.prototype.onMouseDown = function(e, pos) {
            if (onMouseDown) {
                onMouseDown.call(this, e, pos);
            }
            
            // 处理图表点击
            if (this.lossChart && this.lossChart.handleMouseClick(pos)) {
                this.setDirtyCanvas(true, false);
            }
        };
        
        // 扩展updateProgress方法
        const originalUpdateProgress = nodeType.prototype.updateProgress;
        nodeType.prototype.updateProgress = function(data) {
            if (originalUpdateProgress) {
                originalUpdateProgress.call(this, data);
            }
            
            // 更新Loss图表
            if (this.lossChart && data.loss !== undefined && data.loss !== null) {
                this.lossChart.updateLoss(data.current_step, data.loss);
            }
        };
    }
});

// 导出TrainingLossChart类供其他模块使用
window.TrainingLossChart = TrainingLossChart; 